from yolo import YOLO, detect_video
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import cv2
import glob

def detect_img(yolo):
    url = r'D:\桌面\test\1'
    urls = glob.glob(url+'\\*.jpg')
    count = 0
    fps = 8
    out = cv2.VideoWriter('demo.avi', cv2.VideoWriter_fourcc(*'DIVX'), fps, (3840, 2160))
    print(len(urls))
    while True:
        img = urls[count]
        try:
            image = Image.open(img)
        except:
            print('Open Error! Try again!')
            continue
        else:
            r_image = yolo.detect_image(image)
            npimage = np.asarray(r_image)
            out.write(npimage)
        if count > len(urls)-1:
            out.release()
            break
        count += 1

if __name__ == '__main__':
    detect_img(YOLO())